import numpy as np
import torch
import matplotlib.pyplot as plt

def visualize_pca_features_bias(pca_features, labeled_nodes, train_nid, val_nid, test_nid, checkpt_folder, n_components_to_plot=5):
    # 1. 转 numpy, 确保维度
    if isinstance(pca_features, torch.Tensor):
        pca_np = pca_features.detach().cpu().numpy()
    else:
        pca_np = np.array(pca_features)

    num_rows, num_comps = pca_np.shape
    n_plot = min(n_components_to_plot, num_comps)

    labeled_nodes = np.array(labeled_nodes, dtype=np.int64)
    train_nid = np.array(train_nid, dtype=np.int64)
    val_nid = np.array(val_nid, dtype=np.int64)
    test_nid = np.array(test_nid, dtype=np.int64)

    # 2. 将节点 id 映射为 pca_features 的行索引（位置索引）
    def ids_to_positions(query_ids):
        mask = np.isin(labeled_nodes, query_ids)
        positions = np.where(mask)[0]
        return positions

    train_pos = ids_to_positions(train_nid)
    val_pos = ids_to_positions(val_nid)
    test_pos = ids_to_positions(test_nid)

    sets = {
        'All': np.arange(len(pca_np)),
        'Train': train_pos,
        'Val': val_pos,
        'Test': test_pos
    }

    # 3. 计算每个集合在前 n_plot 个主成分上的均值
    # means = {name: [pca_np[idx, i].mean() if len(idx) > 0 else 0 for i in range(n_plot)]
    #          for name, idx in sets.items()}
    # 计算均值和标准差
    means = {name: [pca_np[idx, i].mean() if len(idx) > 0 else 0 for i in range(n_plot)]
             for name, idx in sets.items()}
    stds = {name: [pca_np[idx, i].std() if len(idx) > 0 else 0 for i in range(n_plot)]
            for name, idx in sets.items()}

    # 绘制分组柱状图（带误差棒）
    labels = [f'C{i + 1}' for i in range(n_plot)]
    x = np.arange(n_plot)
    width = 0.2

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.bar(x - 1.5 * width, means['All'], width, yerr=stds['All'], label='All', capsize=5)
    ax.bar(x - 0.5 * width, means['Train'], width, yerr=stds['Train'], label='Train', capsize=5)
    ax.bar(x + 0.5 * width, means['Val'], width, yerr=stds['Val'], label='Val', capsize=5)
    ax.bar(x + 1.5 * width, means['Test'], width, yerr=stds['Test'], label='Test', capsize=5)

    ax.set_xlabel("Principal Components")
    ax.set_ylabel("Mean Component Value")
    ax.set_title("PCA Feature Bias (Top 5 Components)")
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.legend()

    plt.tight_layout()
    plt.savefig(f"{checkpt_folder}/pca_components_grouped_bar.png", dpi=300, bbox_inches='tight')
    plt.close()